package com.feriki.msync.process.source;

import com.alibaba.fastjson.JSONObject;
import com.feriki.msync.constants.ConnectType;
import com.feriki.msync.meta.MilvusCollectionMeta;
import com.feriki.msync.pool.MilvusConnectionPool;
import com.feriki.msync.process.sink.SinkHandler;
import com.feriki.msync.utils.MilvusOperatorUtil;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.orm.iterator.QueryIterator;
import io.milvus.param.R;
import io.milvus.param.collection.FieldType;
import io.milvus.param.dml.QueryIteratorParam;
import io.milvus.response.DescCollResponseWrapper;
import io.milvus.response.QueryResultsWrapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @author luozhida
 * @create 2024-07-04
 */
@Slf4j
public class MilvusSourceHandler extends SourceHandler{

    private MilvusConnectionPool sourcePool;
    private MilvusServiceClient sourceMilvusClient;
    private String milvusSourceSearchExpr;

    public MilvusSourceHandler(MilvusConnectionPool sourcePool,String collectionName,
                               String milvusSourceSearchExpr, long batchSize, long limit) {
        super(collectionName, ConnectType.MILVUS,batchSize,limit);
        this.sourcePool = sourcePool;
        this.milvusSourceSearchExpr = milvusSourceSearchExpr;
        try {
            sourceMilvusClient = sourcePool.getConnection();
        } catch (Exception e) {
            throw new IllegalStateException("源milvus连接获取不到");
        }
    }

    @Override
    public Pair<Boolean,Long> readAndHandle(SinkHandler sinkHandler) {
        try {
            boolean exists = MilvusOperatorUtil.checkCollectionExist(getName(), sourceMilvusClient);
            if (exists) {
                DescCollResponseWrapper sourceDesc = MilvusOperatorUtil.getCollectionDescription(getName(), sourceMilvusClient);
                if (sourceDesc == null) {
                    log.error(String.format("获取不到源collection【%s】的信息", getName()));
                    return Pair.of(false,null);
                }

                MilvusCollectionMeta collectionMeta = new MilvusCollectionMeta();
                collectionMeta.setCollectionName(getName());
                collectionMeta.setShardsNum(sourceDesc.getShardNumber());
                collectionMeta.setCollectionDescription(sourceDesc.getCollectionDescription());
                collectionMeta.setEnableDynamicField(sourceDesc.isDynamicFieldEnabled());
                collectionMeta.getFieldTypesFrom(sourceDesc.getFields(),sourceDesc.getVectorField().getName());
                if(!sinkHandler.prepare(getName(),collectionMeta)){
                    return Pair.of(false,null);
                }
                List<String> fields = sourceDesc.getFields().stream().map(FieldType::getName).collect(Collectors.toList());
                //迭代查询
                QueryIteratorParam.Builder requestParamBuilder = QueryIteratorParam.newBuilder()
                        .withBatchSize(getBatchSize())
                        .withLimit(getLimit()<0?-1:getLimit())
                        .withCollectionName(getName())
                        .withOutFields(fields);
                if (StringUtils.isBlank(milvusSourceSearchExpr)) {
                    //使用不会命中的条件进行查询全部
                    String primaryKey = sourceDesc.getPrimaryField().getName();
                    DataType primaryDataType = sourceDesc.getPrimaryField().getDataType();

                    if (primaryDataType.getNumber() == DataType.String.getNumber()
                            || primaryDataType.getNumber() == DataType.VarChar.getNumber()) {
                        requestParamBuilder.withExpr(primaryKey + " != '-99999'");
                    } else if (primaryDataType.getNumber() == DataType.Int64.getNumber()) {
                        requestParamBuilder.withExpr(primaryKey + " != -99999");
                    } else {
                        log.error(String.format("源collection【%s】的主键字段类型非Int64或VarChar", getName()));
                        return Pair.of(false,null);
                    }
                }else{
                    //使用用户指定的条件进行查询
                    requestParamBuilder.withExpr(milvusSourceSearchExpr);
                }
                long totalSuccess = 0;

                QueryIteratorParam requestParam = requestParamBuilder.build();
                log.info("即将对源collection【{}】执行查询条件【{}】{}",getName(),requestParam.getExpr(),requestParam.getLimit()<0?"":"查询"+requestParam.getLimit()+"条");
                R<QueryIterator> queryIteratorR = sourceMilvusClient.queryIterator(requestParam);
                if (queryIteratorR.getStatus() == R.Status.Success.getCode() && queryIteratorR.getData() != null) {
                    List<QueryResultsWrapper.RowRecord> rowRecords;
                    while ((rowRecords = queryIteratorR.getData().next()) != null && !rowRecords.isEmpty()) {
                        List<JSONObject> jsonObjects = rowRecords.stream()
                                .map(QueryResultsWrapper.RowRecord::getFieldValues)
                                .map(v ->
                                        (JSONObject) JSONObject.toJSON(v)
                                ).collect(Collectors.toList());

                        long success = sinkHandler.handler(getName(),jsonObjects);
                        totalSuccess = totalSuccess + success;
                    }
                    queryIteratorR.getData().close();
                }else{
                    log.error("源【{}】查询迭代器失败:{}", getName(),queryIteratorR.getMessage());
                }
                return Pair.of(true,totalSuccess);
            } else {
                //skip
                log.error(String.format("源【%s】不存在!", getName()));
            }
        } catch (Throwable e) {
            log.error("从源【{}】同步记录到目标【{}】出现异常", getName(), sinkHandler.getName(),e);
        } finally {
            this.destroy();
            sinkHandler.destroy();
        }
        return Pair.of(false,null);
    }

    @Override
    public void destroy() {
        if (sourceMilvusClient != null) {
            sourcePool.releaseConnection(sourceMilvusClient);
            sourceMilvusClient = null;
        }
    }
}
